import torch
import argparse


def parse_argument():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    ## dir params
    parser.add_argument('--epoch_1_path', type=str,
                        default='checkpoints/epoch=002-step=600000-val_AP=0.33.ckpt')
    parser.add_argument('--epoch_2_path', type=str,
                        default='RVT/cop3adl1/checkpoints/7.17new.ckpt')

    parser.add_argument('--save_path', type=str,
                        default='checkpoints/Res_pretrain2.ckpt')

    return parser


if __name__ == '__main__':
    args, _ = parse_argument().parse_known_args(None)

    pretrained_ckpt = torch.load(args.epoch_1_path, map_location='cuda:0')
    target_ckpt = torch.load(args.epoch_2_path, map_location='cuda:0')
    # print(pretrained_ckpt.keys())


    pretrained_state_dict = pretrained_ckpt['state_dict']
    print(pretrained_state_dict.keys())
    target_state_dict = target_ckpt['state_dict']
    print(target_state_dict.keys())
    assert target_state_dict.keys() == pretrained_state_dict.keys()

    for k in target_state_dict.keys():

        if target_state_dict[k].shape == pretrained_state_dict[k].shape:
            # print(k)
            target_state_dict[k] = pretrained_state_dict[k]
        else:
            print(k)


    save_key = {'state_dict': target_state_dict}

    torch.save(save_key, args.save_path)







